import os
import dashscope
import bs4
from langchain_community.vectorstores import Chroma
from langchain_core.runnables import RunnableLambda
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.documents import Document
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.embeddings import DashScopeEmbeddings
from langchain.text_splitter import CharacterTextSplitter, RecursiveCharacterTextSplitter

# 获取环境变量
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
DASHSCOPE_API_KEY = os.getenv("DASHSCOPE_API_KEY")

dash_embeddings = DashScopeEmbeddings(dashscope_api_key=DASHSCOPE_API_KEY,
                                      model="text-embedding-v3")
# 文档模拟数据
documents = [
    Document(page_content="猫是懒散可爱的动物",
              metadata={"source":"动物宠物文档"}
    )
]
# 使用bs4库进行网页数据提取
web_loader = WebBaseLoader(
    web_path=["https://www.news.cn/fortune/20250612/abb442b0ba95444eb9ffa17ef3d25230/c.html"],
    bs_kwargs=dict(
        parse_only = bs4.SoupStrainer(class_=("main-left","title"))
    )
)
#定义文档分割器
splitter = RecursiveCharacterTextSplitter(chunk_size=100,chunk_overlap=20)
web_docs = splitter.split_documents(web_loader.load())

# 实例化向量空间
print("创建向量数据库...")
vector_database = Chroma.from_documents(
    documents=documents + web_docs,
    embedding=dash_embeddings,
    persist_directory="./chroma_db"  # 添加持久化目录
)
print("向量数据库创建成功")


# 生成可运行检测器，并包装成langchain的Runnable接口，可以与其他组件链式结合
def search_docs(input_dict):
    query = input_dict.get("input", "")
    return vector_database.similarity_search(query, k=2)

doc_retriever = RunnableLambda(search_docs)

